package com.chatplus.application.processor.channel.chat.xunfei;

import baidumodel.constant.ChatRoleConstant;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.domain.dto.ChatConfigDto;
import com.chatplus.application.domain.entity.basedata.ApiKeyEntity;
import com.chatplus.application.processor.AIChannelServiceProvider;
import com.chatplus.application.processor.channel.chat.IChatAiProcessorService;
import com.chatplus.application.processor.channel.chat.xunfei.dto.XunFeiMsgDTO;
import com.chatplus.application.processor.channel.chat.xunfei.dto.XunFeiRequestDTO;
import com.chatplus.application.processor.constant.AIConstants;
import com.chatplus.application.processor.constant.AiPlatformConstants;
import com.chatplus.application.processor.handler.dto.ChatRecordMessage;
import com.chatplus.application.processor.listener.XunFeiEventSourceListener;
import com.chatplus.application.service.basedata.ApiKeyService;
import com.chatplus.application.service.basedata.ConfigService;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import okhttp3.HttpUrl;
import okhttp3.OkHttpClient;
import okhttp3.Request;
import okhttp3.WebSocket;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Service;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.text.SimpleDateFormat;
import java.util.ArrayList;
import java.util.Base64;
import java.util.Date;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
import java.util.TimeZone;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;


/**
 * 讯飞 星火 机器人处理器
 *
 * @author chj
 * @date 2023/12/27
 **/
@Service(value = AIChannelServiceProvider.SERVICE_NAME_PRE + "XunFei")
public class XunFeiXhProcessor extends IChatAiProcessorService {
    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(XunFeiXhProcessor.class);
    private final ConfigService configService;
    private final ApiKeyService apiKeyService;

    private final ObjectMapper objectMapper;

    public XunFeiXhProcessor(ConfigService configService
            , ApiKeyService apiKeyService
            , ObjectMapper objectMapper) {
        this.configService = configService;
        this.apiKeyService = apiKeyService;
        this.objectMapper = objectMapper;
    }

    private Map<String, String> model2URL = Map.of(
            "general", "v1.1",
            "generalv2", "v2.1",
            "generalv3", "v3.1"
    );

    @Override
    public void process(List<ChatRecordMessage> messageContext) throws JsonProcessingException {
        ChatConfigDto chatConfigDto = configService.getChatConfig();
        ChatConfigDto.XunFeiBean xunFeiBean = chatConfigDto.getXunFei();
        if (xunFeiBean == null) {
            replyErrorMessage("加载系统配置失败，连接已关闭！！！");
            return;
        }
        List<ApiKeyEntity> apiKeys = apiKeyService.getApiKeys(getChannel(), "chat");
        if (CollectionUtils.isEmpty(apiKeys)) {
            replyErrorMessage("抱歉😔😔😔，系统已经没有可用的 API KEY，请联系管理员！");
            replyErrorMessage(AIConstants.ERROR_MSG);
            replyErrorMessage(AIConstants.ERR_IMG);
            return;
        }
        ApiKeyEntity apiKeyEntity = apiKeys.get(0);
        // 需要用竖线（|）将 APPID, APIKey 和 APISecret 按照顺序串接起来填入
        String appId = apiKeyEntity.getValue().split("\\|")[0];
        String apiKey = apiKeyEntity.getValue().split("\\|")[1];
        String secretKey = apiKeyEntity.getValue().split("\\|")[2];
        String modelName = getSession().getAttributes().get("modelName").toString();
        String apiUrl = xunFeiBean.getApiUrl();
        apiUrl = apiUrl.replace("{version}", model2URL.get(modelName));
        xunFeiBean.setApiUrl(apiUrl);
        // 生成鉴权访问大模型的路径
        String authUrl = getAuthUrl(apiUrl, apiKey, secretKey);
        // 鉴权方法生成失败，直接返回 null
        if (authUrl == null) {
            return;
        }
        OkHttpClient okHttpClient = new OkHttpClient.Builder().build();
        // 将 https/http 连接替换为 ws/wss 连接
        String url = authUrl.replace("http://", "ws://").replace("https://", "wss://");
        Request request = new Request.Builder().url(url).build();
        String prompt = messageContext.stream()
                .filter(msg -> msg.getReply() == null)
                .map(ChatRecordMessage::getPrompt)
                .collect(Collectors.joining());
        // 建立 wss 连接
        WebSocket webSocket = okHttpClient.newWebSocket(request, new XunFeiEventSourceListener(getSession(), getChannel(), prompt));
        // 组装请求参数
        List<XunFeiMsgDTO> messages = new ArrayList<>();
        AtomicInteger index = new AtomicInteger();
        messageContext.forEach(msg -> {
            XunFeiMsgDTO promptMessage = new XunFeiMsgDTO(ChatRoleConstant.USER, msg.getPrompt(), index.get());
            index.getAndIncrement();
            messages.add(promptMessage);
            if (StringUtils.isNotEmpty(msg.getReply())) {
                XunFeiMsgDTO replyMessage = new XunFeiMsgDTO(ChatRoleConstant.ASSISTANT, msg.getPrompt(), index.get());
                messages.add(replyMessage);
            }
        });
        XunFeiRequestDTO requestDTO = getRequestParam(messages, xunFeiBean, modelName, appId);
        // 发送请求
        webSocket.send(objectMapper.writeValueAsString(requestDTO));
    }

    @Override
    public String getChannel() {
        return AiPlatformConstants.XUN_FEI;
    }

    /**
     * 生成鉴权方法，具体实现不用关心，这是讯飞官方定义的鉴权方式
     *
     * @return 鉴权访问大模型的路径
     */
    public String getAuthUrl(String authUrl, String apiKey, String apiSecret) {
        try {
            authUrl = authUrl.replace("ws://", "http://").replace("wss://", "https://");
            URL url = new URL(authUrl);
            // 时间
            SimpleDateFormat format = new SimpleDateFormat("EEE, dd MMM yyyy HH:mm:ss z", Locale.US);
            format.setTimeZone(TimeZone.getTimeZone("GMT"));
            String date = format.format(new Date());
            // 拼接
            String preStr = "host: " + url.getHost() + "\n" +
                    "date: " + date + "\n" +
                    "GET " + url.getPath() + " HTTP/1.1";
            // SHA256加密
            Mac mac = Mac.getInstance("hmacsha256");
            SecretKeySpec spec = new SecretKeySpec(apiSecret.getBytes(StandardCharsets.UTF_8), "hmacsha256");
            mac.init(spec);
            byte[] hexDigits = mac.doFinal(preStr.getBytes(StandardCharsets.UTF_8));
            // Base64加密
            String sha = Base64.getEncoder().encodeToString(hexDigits);
            // 拼接
            String authorizationOrigin = String.format("api_key=\"%s\", algorithm=\"%s\", headers=\"%s\", signature=\"%s\"", apiKey, "hmac-sha256", "host date request-line", sha);
            // 拼接地址
            HttpUrl httpUrl = Objects.requireNonNull(HttpUrl.parse("https://" + url.getHost() + url.getPath())).newBuilder().
                    addQueryParameter("authorization", Base64.getEncoder().encodeToString(authorizationOrigin.getBytes(StandardCharsets.UTF_8))).
                    addQueryParameter("date", date).
                    addQueryParameter("host", url.getHost()).
                    build();
            return httpUrl.toString();
        } catch (Exception e) {
            LOGGER.message("鉴权方法中发生错误").exception(e).error();
            return null;
        }
    }

    /**
     * 获取请求参数
     *
     * @param msgList 发送给大模型的消息，可以包含上下文内容
     * @return 请求DTO，该 DTO 转 json 字符串后生成的格式参考 resources/demo-json/domain.json
     */
    public XunFeiRequestDTO getRequestParam(List<XunFeiMsgDTO> msgList, ChatConfigDto.XunFeiBean xunFeiBean, String modelName, String appId) {
        XunFeiRequestDTO requestDTO = new XunFeiRequestDTO();
        requestDTO.setHeader(new XunFeiRequestDTO.HeaderDTO(appId, getSession().getAttributes().get("userId").toString()));
        requestDTO.setParameter(new XunFeiRequestDTO.ParameterDTO(new XunFeiRequestDTO.ParameterDTO.ChatDTO(modelName
                , Float.parseFloat(xunFeiBean.getTemperature() + ""), xunFeiBean.getMaxTokens())));
        requestDTO.setPayload(new XunFeiRequestDTO.PayloadDTO(new XunFeiRequestDTO.PayloadDTO.MessageDTO(msgList)));
        return requestDTO;
    }
}
